from fastapi import FastAPI, WebSocket
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from env import World
from agent import Agent
import json
import numpy as np
import random

app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")

# Serve index.html at root route
@app.get("/")
async def get():
    return FileResponse('static/index.html')

class GameServer:
    def __init__(self):
        self.world = World()
        self.agent = Agent()
        self.current_maze = None
        
    def generate_maze(self, size, complexity):
        """Generate maze using recursive backtracking"""
        maze = np.ones((size, size))
        start = (1, 1)
        self._carve_maze(maze, start[0], start[1])
        
        # Add some random walls based on complexity
        for _ in range(int(size * size * complexity)):
            x = random.randint(1, size-2)
            y = random.randint(1, size-2)
            maze[x, y] = 1
            
        self.current_maze = maze.tolist()
        return self.current_maze

    async def game_loop(self, websocket: WebSocket):
        await websocket.accept()
        
        while True:
            # Get world state
            state = {
                "agent": {
                    "position": self.agent.position,
                    "energy": self.agent.energy,
                    "score": self.agent.score
                },
                "world": {
                    "terrain": self.world.terrain,
                    "platforms": self.world.platforms,
                    "obstacles": self.world.obstacles,
                    "maze": self.current_maze
                }
            }
            
            # Send state to client
            await websocket.send_json(state)
            
            # Get client input if any
            try:
                data = await websocket.receive_json()
                # Handle client commands
                if data.get("command") == "reset":
                    self.world = World()
                    self.agent = Agent()
                elif data.get("command") == "generate_maze":
                    size = data.get("size", 10)
                    complexity = data.get("complexity", 0.3)
                    self.generate_maze(size, complexity)
            except:
                break

game = GameServer()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await game.game_loop(websocket) 